import numpy as np
import math
import plotly.graph_objects as go
from plotly.subplots import make_subplots

n = 32
k = 1
p=0.25
def have_at_least_x_slots(p, k=1):
    p_sum = 0 
    for i in range(k, n+1):
        binomial_coefficient = math.factorial(n) / (math.factorial(i) * math.factorial(n-i))
        p_i = binomial_coefficient * (p**i) * ((1-p)**(n-i))
        p_sum += p_i
    return p_sum

def cum_prob_dist(p):
    p_cum_sum = []
    k=0
    for i in range(k, n+1):
        binomial_coefficient = math.factorial(n) / (math.factorial(i) * math.factorial(n-i))
        p_i = binomial_coefficient * (p**i) * ((1-p)**(n-i))
        p_cum_sum.append(p_i)
    return p_cum_sum

def draw_probability(_y):
    fig = make_subplots(specs=[[{"secondary_y": True}]])
    for ix, (y, name) in enumerate(_y):
        dash = "dot" if ix > 1 else "dash" if ix == 1 else "solid"
        x = list(range(0, len(y)))
        fig.add_trace(
            go.Scatter(x=x, y=y, mode='lines', name=f"prob. dist. {name}", line = dict(width=3, dash=dash)),  
            secondary_y=False
        )
        fig.add_trace(
            go.Scatter(x=x, y=np.cumsum(y), mode='lines', name=f"cum. prob. dist. {name}", line = dict(width=3, dash=dash)), 
            secondary_y=True
        )
    fig.update_layout(
        xaxis_title='Slots',
        legend=dict(x=0.7, y=0.9),
        width = 900,
        height = 350,
        margin=dict(l=0, r=0, t=0, b=0),
        xaxis=dict(
            showgrid=True,
            gridwidth=1,
            gridcolor='white',
            title_font=dict(size=18)
        ),
        yaxis=dict(
            title='Probability', 
            color='blue',
            showgrid=True,
            gridwidth=1,
            gridcolor='white',
            title_font=dict(size=18)
        ),
        yaxis2=dict(
            title='Cumulative Probability', 
            color='red', 
            overlaying='y', 
            side='right',
            showgrid=True,
            gridwidth=1,
            gridcolor='white',
            title_font=dict(size=18)
        ),
        font=dict(
            size=16
        ),
    )
                      
    return fig

y_25 = cum_prob_dist(0.25)
y_10 = cum_prob_dist(0.10)
y_50 = cum_prob_dist(0.50)
fig = draw_probability([(y_25, "25%"), (y_10, "10%"), (y_50, "50%")])
fig.write_image("probability_dist.png")
fig.show()

games = list(range(1, 21)) 
p = .25

def continue_biasing(p):
    probabilities = [] 
    for n in games:
        k = 0
        comb = math.comb(n, k)
        P_X_equals_0 = comb * (p**k) * ((1-p)**(n-k))
        P_X_geq_1 = 1 - P_X_equals_0

        probabilities.append(P_X_geq_1)
    return probabilities

probabilities_13 = continue_biasing(have_at_least_x_slots(p, 13)) # get 12 slots
probabilities_12 = continue_biasing(have_at_least_x_slots(p, 12)) # get 12 slots
probabilities_11 = continue_biasing(have_at_least_x_slots(p, 11)) # get 11 slots
probabilities_10 = continue_biasing(have_at_least_x_slots(p, 10)) # get 10 slots
probabilities_09 = continue_biasing(have_at_least_x_slots(p, 9))  # get 9 slots

fig = go.Figure()
for probabilities, name in zip([probabilities_13, 
                                probabilities_12, 
                                probabilities_11, 
                                probabilities_10, 
                                probabilities_09
                               ], 
                               ["getting 13 slots", 
                                "getting 12 slots", 
                                "getting 11 slots", 
                                "getting 10 slots", 
                                "getting 9 slots"
                               ]):
    fig.add_trace(
        go.Scatter(x=games, y=probabilities, mode='lines+markers', name=name)
    )

for x, sl in zip([2, 4, 8, 16],[1,2,3,4]):
    fig.add_shape(
        type='line',
        x0=x, y0=0, x1=x, y1=1,
        line=dict(color='Red',dash="dash"),
    )
    fig.add_annotation(x=x+0.95, y=0.05,
                       text=f"Controlling<br>{sl} tail slots",
                       showarrow=False,
                       font=dict(size=11, color="black"),
                       bgcolor="white",
                       opacity=0.7,
                      )

fig.update_layout(
    title='Probability of getting X slots with 25% validator share',
    xaxis=dict(
        title='Number of Possibilities (randao_reveal vs. missed slot)', 
        titlefont=dict(size=18),
        tickvals=list(range(0,21, 2))
    ),
    yaxis=dict(title='Probability', titlefont=dict(size=18)),
    showlegend=True,
    width = 900,
    height = 400,
    margin=dict(l=0, r=0, t=50, b=0),
    legend=dict(font=dict(size=16)),
    font=dict(
        size=16
    ),
)
fig.write_image("probabilities.png")
fig.show()